package org.san;

import org.apache.zookeeper.CreateMode;
import org.apache.zookeeper.KeeperException;
import org.apache.zookeeper.ZooDefs;
import org.apache.zookeeper.ZooKeeper;
import org.apache.zookeeper.data.Stat;

import java.util.List;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.CountDownLatch;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.atomic.AtomicInteger;

/**
 * @author gaodong
 * @version 1.0
 * @date 2021/6/10 14:07
 */
public class ZooKeeperReenterantLock {

    private ZooKeeper zooKeeper;
    private String root;
    private String lockPathPrefix;
    private final static String prefix = "/id_";

    private ConcurrentHashMap<Thread, LockData> threadDataMap = new ConcurrentHashMap<>();

    private final class LockData {
        Thread owner;
        String lockPath;
        AtomicInteger lockCount = new AtomicInteger(1);

        public LockData(Thread thread, String path) {
            owner = thread;
            lockPath = path;
        }
    }

    // region 外部调用接口

    public ZooKeeperReenterantLock(ZooKeeper zooKeeper, String root) {
        this.zooKeeper = zooKeeper;
        this.root = root;
        this.lockPathPrefix = root + prefix;

        try {
            Stat stat = zooKeeper.exists(root, null);
            if (stat == null) {
                zooKeeper.create(root, dataBytes, ZooDefs.Ids.OPEN_ACL_UNSAFE, CreateMode.PERSISTENT);
            }
        } catch (Exception e) {
        }
    }

    /**
     * 获取锁
     *
     * @return
     * @throws Exception
     */
    public boolean acquire() throws Exception {
        Thread currentThread = Thread.currentThread();
        LockData lockData = threadDataMap.get(currentThread);
        if (lockData != null) {
            lockData.lockCount.incrementAndGet();
            return true;
        }

        String lockPath = createLock();
        checkLock(lockPath, -1);

        if (lockPath != null) {
            threadDataMap.put(currentThread, new LockData(currentThread, lockPath));
            return true;
        }

        return false;
    }

    /**
     * 获取锁，指定之间内未成功，自动释放
     *
     * @param timeout
     * @param unit
     * @return
     */
    public boolean acquire(long timeout, TimeUnit unit) throws Exception {
        Thread currentThread = Thread.currentThread();
        LockData lockData = threadDataMap.get(currentThread);
        if (lockData != null) {
            lockData.lockCount.incrementAndGet();
            return true;
        }

        String lockPath = createLock();
        if (!checkLock(lockPath, unit.toNanos(timeout))) {
            zooKeeper.delete(lockPath, -1);
            return false;
        }

        if (lockPath != null) {
            threadDataMap.put(currentThread, new LockData(currentThread, lockPath));
            return true;
        }

        return false;
    }

    /**
     * 释放锁
     *
     * @throws Exception
     */
    public void release() throws Exception {
        Thread currentThread = Thread.currentThread();
        LockData lockData = threadDataMap.get(currentThread);
        if (lockData == null) {
            throw new IllegalMonitorStateException("You do not own the lock: " + root);
        }

        int newLockCount = lockData.lockCount.decrementAndGet();

        if (newLockCount > 0) {
            return;
        }

        if (newLockCount < 0) {
            throw new IllegalMonitorStateException("Lock count has gone negative for lock: " + root);
        }

        try {
            zooKeeper.delete(lockData.lockPath, -1);
        } finally {
            threadDataMap.remove(currentThread);
        }
    }

    // endregion

    // region 内部方法

    private final byte[] dataBytes = "".getBytes();

    private String createLock() throws KeeperException, InterruptedException {
        return zooKeeper.create(lockPathPrefix, dataBytes, ZooDefs.Ids.OPEN_ACL_UNSAFE, CreateMode.EPHEMERAL_SEQUENTIAL);
    }

    private boolean checkLock(String path, long timeoutNano) throws KeeperException, InterruptedException {
        long start = System.nanoTime();
        do {
            List<String> children = getSortedChildren();
            int index = children.indexOf(path.substring(path.lastIndexOf("/") + 1));
            if (index == 0) {
                return true;
            }

            long usedNano = System.nanoTime() - start;
            if (timeoutNano != -1) {
                timeoutNano -= usedNano;
                if (timeoutNano <= 0) {
                    return false;
                }
                start = System.nanoTime();
            }

            watch(lockPathPrefix + children.get(index - 1), timeoutNano);

        } while (true);
    }

    private List<String> getSortedChildren() throws KeeperException, InterruptedException {
        List<String> children = zooKeeper.getChildren(root, null);
        children.sort((a, b) -> {
            int na = Integer.parseInt(a.substring(3)), nb = Integer.parseInt(b.substring(3));
            // na ^ nb judge the symbol of na/nb
            // don't use (nb-na) in case of overflow when na ^ nb < 0
            return (na ^ nb) >= 0 ? (na - nb) : (na > nb ? -1 : 1);
        });
        return children;
    }

    private void watch(String path, long timeoutNano) throws KeeperException, InterruptedException {
        CountDownLatch countDownLatch = new CountDownLatch(1);
        Stat stat = zooKeeper.exists(path, event -> countDownLatch.countDown());
        if (stat != null) {
            if (timeoutNano == -1) {
                countDownLatch.wait();
            } else {
                countDownLatch.await(timeoutNano, TimeUnit.NANOSECONDS);
            }
        }
    }

    // endregion
}
